import os
import torch
import torchvision.models as models
import torchvision
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
import torch.utils.data as torchdata
from torch.optim import lr_scheduler
from dataset.facedata import LFW_np,LFW_zip
from models.spherenet import *
from dataset.casia_sphere import ImageDataset,dataset_load
from utils.train import train,trainlog
from model_prune.filter_pruner_spherenet import FilterPruner
from model_prune.prune_utils import *

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
def refine_preprune(prepruned_index,size):
    refined_index = []
    inter = size[0] * size[1]
    for idx in prepruned_index:
        refined_index = refined_index + range(idx*inter, (idx+1)*inter)
    return refined_index


def total_num_filters(model):
    filters = 0
    for name1, module1 in model._modules.items():
        if isinstance(module1, torch.nn.modules.conv.Conv2d):
            filters = filters + module1.out_channels
        for name2, module2 in module1._modules.items():
            if isinstance(module2, torch.nn.modules.conv.Conv2d):
                filters = filters + module2.out_channels
            for name3, module3 in module2._modules.items():
                if isinstance(module3, torch.nn.modules.conv.Conv2d):
                    filters = filters + module3.out_channels
                for name4, module4 in module3._modules.items():
                    if isinstance(module4, torch.nn.modules.conv.Conv2d):
                        filters = filters + module4.out_channels
    return filters

class SpherenetFilterPruner(FilterPruner):
    # Inherited a FilterPruner class

    # define a custom forward function in which hooks are registered to calculate grads
    def forward(self, x):
        self.activations = []
        activation_index = 0
        x = self.model.conv1_1(x)
        self.attach_action_and_hook(x, activation_index, ('conv1_1',))
        activation_index += 1
        x = self.model.relu1_1(x)

        x1 = self.model.conv1_2(x)
        self.attach_action_and_hook(x1, activation_index, ('conv1_2',))
        activation_index += 1
        x1 = self.model.relu1_2(x1)
        x1 = self.model.conv1_3(x1)
        x1 = self.model.relu1_3(x1)
        x = x1 + x

        # conv2
        x = self.model.conv2_1(x)
        self.attach_action_and_hook(x, activation_index, ('conv2_1',))
        activation_index += 1
        x = self.model.relu2_1(x)

        x1 = self.model.conv2_2(x)
        self.attach_action_and_hook(x1, activation_index, ('conv2_2',))
        activation_index += 1
        x1 = self.model.relu2_2(x1)
        x1 = self.model.conv2_3(x1)
        x1 = self.model.relu2_3(x1)
        x = x1 + x

        x1 = self.model.conv2_4(x)
        self.attach_action_and_hook(x1, activation_index, ('conv2_4',))
        activation_index += 1
        x1 = self.model.relu2_4(x1)
        x1 = self.model.conv2_5(x1)
        x1 = self.model.relu2_5(x1)
        x = x1 + x

        # conv3
        x = self.model.conv3_1(x)
        self.attach_action_and_hook(x, activation_index, ('conv3_1',))
        activation_index += 1
        x = self.model.relu3_1(x)

        x1 = self.model.conv3_2(x)
        self.attach_action_and_hook(x1, activation_index, ('conv3_2',))
        activation_index += 1
        x1 = self.model.relu3_2(x1)
        x1 = self.model.conv3_3(x1)
        x1 = self.model.relu3_3(x1)
        x = x1 + x

        x1 = self.model.conv3_4(x)
        self.attach_action_and_hook(x1, activation_index, ('conv3_4',))
        activation_index += 1
        x1 = self.model.relu3_4(x1)
        x1 = self.model.conv3_5(x1)
        x1 = self.model.relu3_5(x1)
        x = x1 + x

        x1 = self.model.conv3_6(x)
        self.attach_action_and_hook(x1, activation_index, ('conv3_6',))
        activation_index += 1
        x1 = self.model.relu3_6(x1)
        x1 = self.model.conv3_7(x1)
        x1 = self.model.relu3_7(x1)
        x = x1 + x

        x1 = self.model.conv3_8(x)
        self.attach_action_and_hook(x1, activation_index, ('conv3_8',))
        activation_index += 1
        x1 = self.model.relu3_8(x1)
        x1 = self.model.conv3_9(x1)
        x1 = self.model.relu3_9(x1)
        x = x1 + x

        x = self.model.conv4_1(x)
        self.attach_action_and_hook(x, activation_index, ('conv4_1',))
        activation_index += 1
        x = self.model.relu4_1(x)

        x1 = self.model.conv4_2(x)
        self.attach_action_and_hook(x1, activation_index, ('conv4_2',))
        activation_index += 1
        x1 = self.model.relu4_2(x1)
        x1 = self.model.conv4_3(x1)
        x1 = self.model.relu4_3(x1)
        x = x1 + x

        x = x.view(x.size(0),-1)
        x = self.model.fc5(x)
        if self.model.feature: return x

        x = self.model.fc6(x)
        return x
    # define a custom prune function to prune filters (functions in prune_utils may be helpful)
    def prune_conv_layer(self,model, prune_target):
        channels_now = 3
        prepruned_index = []
        repeat_index = []
        for i, (name1, module1) in enumerate(model._modules.items()):
            if isinstance(module1, nn.Conv2d):
                if (name1, ) in prune_target.keys():
                    prune_index =  prune_target[(name1,)]
                else:
                    prune_index = []
                if name1 in ['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1']:
                    repeat_index = prune_index

                elif name1 in ['conv1_3','conv2_3', 'conv2_5','conv3_3', 'conv3_5',
                             'conv3_7','conv3_9','conv4_3']:
                    prune_index = repeat_index

                module1 = prune_conv_weight(module1, prepruned_index, prune_index)
                prepruned_index = prune_index  # prepruned index <- prune index
                channel_now = module1.out_channels  # channel_now <- out_channels
            if isinstance(module1, nn.PReLU):
                module1 = prune_prelu_weight(module1, prepruned_index)

            if name1 == 'fc5':
                if (name1, ) in prune_target.keys():
                    prune_index =  prune_target[(name1,)]
                else:
                    prune_index = []

                prepruned_index = refine_preprune(prepruned_index,(7,6))
                module1 = prune_fc_weight(module1, prepruned_index,prune_index)

        return model


if __name__ == '__main__':
    # data
    data_set = {}
    data_loader = {}
    data_set['val'] = LFW_zip(lfw_path='/media/hszc/data/lfw.zip',
                              pairtxt='/home/hszc/zhangchi/sphereface_pytorch/data/pairs.txt',
                              landmarktxt='/home/hszc/zhangchi/sphereface_pytorch/data/lfw_landmark.txt')

    data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=32, num_workers=4,
                                              shuffle=False, pin_memory=True)


    data_set['webface_zip'] = '/media/hszc/data/webface.zip'
    data_set['casia_landmark'] = '/home/hszc/zhangchi/sphereface_pytorch/data/casia_landmark.txt'

    # data_loader['train'] = ImageDataset(data_set['webface_zip'],dataset_load,
    #                                     data_set['casia_landmark'],
    #                                     name='sphere20a:train',
    #     bs=256,shuffle=True,nthread=6,imagesize=128)

    print 'data prepared'

    criterion = AngleLoss()
    #
    #
    model = sphere20a(classnum=10574,feature=False)
    converted_weights = torch.load('/home/hszc/zhangchi/sphereface_pytorch/model/sphere20a_20171020.pth')
    model.load_state_dict(converted_weights)

    fine_tuner = SpherenetFilterPruner(model=model,
                                       data_set = data_set,
                                       data_loader=data_loader,
                                       criterion=criterion,
                                       useCuda=4)
    fine_tuner.perform_prune(save_dir='/home/hszc/zhangchi/channel-prune/prune_result/sphere2',
                             proportion=0.7,
                             num_filters_per_iter=128,
                             epochs_after_perprune=6,
                             epochs_after_wholeprune=8,
                             bs=256,
                             nthread=6)

